flashatten takeaways EP2 工程

工程相关细节

overview

实际工程实现中牵涉到一些实现细节，主要是针对不同的场景，会有兼容性的设置。体现在代码中的复杂性。
以 flashattention源码为主例，参考剔除了冗余工程实现的相对较为简化的 libflash_atten样例工程记录源码中理解的细节。
flashatten源码完全编译要40G左右内存，建议直接先从最小案例开始，源码理解flashatten已足够。需要编译实验的更多的是CUTE的底层原理，而非flashatten。
涉及较多C++ 和部分CUDA、CUTE 知识
对于其中涉及的底层CUTE实现，因为内容太多，详细的分析不放在这个文档中，分开放在了专门解析CUTE的文档记录中。

下面对工程代码所需理解的若干概念解释，过程中各个概念部分用的例子都相对独立，但是所使用的变量和中间量的名称都尽量与工程代码中保持一致或靠近（除了原理推理）：

kv cache

kvcache 是 K和V两个矩阵的cache，为了方便K，V重复利用。kvcache在如自回归解码时增量预测的情况下发挥作用，此时Q的sequence length远小于K,V的sequence length。其过程如下： Transformer 自回归解码时，每生成一个新 token，需要用当前 query 去关注所有历史 token。若不缓存，每一步都要重算所有历史 token 的 $K$ 和 $V$，复杂度退化为 $O(N^2)$。设已生成长度 $t$ 的序列，下一步计算第 $t+1$ 个 token：

无 KV Cache的情况（重算）：$$ O_t = \operatorname{softmax}\!\left(\frac{Q_t\, K_{1:t}^{\,T}}{\sqrt{d}}\right) V_{1:t},\qquad K_{1:t}=X_{1:t} W_K,\ V_{1:t}=X_{1:t} W_V $$ 每步重算 $K_{1:t}, V_{1:t}$，成本随 $t$ 线性增长，总复杂度 $O(N^2 d)$。
有 KV Cache情况（增量拼接）：缓存历史 $K_{1:t}, V_{1:t}$，仅计算新 token 的 $k_{t+1}, v_{t+1}$：$$ k_{t+1} = x_{t+1} W_K,\qquad v_{t+1} = x_{t+1} W_V $$ $$ K_{1:t+1} = \begin{bmatrix} K_{1:t} \\ k_{t+1} \end{bmatrix},\qquad V_{1:t+1} = \begin{bmatrix} V_{1:t} \\ v_{t+1} \end{bmatrix} $$ $$ \boxed{\;O_{t+1} = \operatorname{softmax}\!\left(\frac{q_{t+1}\, K_{1:t+1}^{\,T}}{\sqrt{d}}\right) V_{1:t+1}\;} $$

kvcache的核心收益为：在无Cache的情况下，每步计算 $K, V$ 复杂度为$O(t \cdot d)$，总解码复杂度为$O(N^2 d)$；有 KV Cache情况下，每步计算 $K, V$ 复杂度为$O(d)$（仅新 token），总解码复杂度为$O(N d)$。

源代码接口

了解CUTE库之后发现，flashatten 的源码与CUTE同样，用到了很多C++模板元编程，为的是提高代码通用性的同时减少编译的负担。最直观的体现是这些平级的文件中的实现：

flash_fwd_kernels/
├── flash_fwd_hdim32_fp16_sm80.cu
├── flash_fwd_hdim64_fp16_sm80.cu
├── flash_fwd_hdim64_e4m3_sm90.cu
├── flash_fwd_hdim64_fp16_sm90.cu
├── flash_fwd_hdim96_fp16_sm80.cu
├── flash_fwd_hdim128_fp16_sm80.cu
├── flash_fwd_hdim128_e4m3_sm90.cu
├── flash_fwd_hdim128_fp16_sm90.cu
├── flash_fwd_hdim160_fp16_sm80.cu
├── flash_fwd_hdim192_fp16_sm80.cu
├── flash_fwd_hdim224_fp16_sm80.cu
├── flash_fwd_hdim256_fp16_sm80.cu
├── flash_fwd_hdim256_e4m3_sm90.cu
└── flash_fwd_hdim256_fp16_sm90.cu

不是每一个文件在各台机器的每一个实现中都会被编译，我的机器不是hopper类型的，计算能力是sm80 的，因此这里以sm80的环境中做的实验和情况举例。

顶层业务代码在调用时会调两种接口，分别对应是否将K、V矩阵分块进行计算的情况，即 run_flash_fwd 接口和 run_flash_splitkv_fwd 接口: 这是 run_flash_fwd 接口：

template<typename Kernel_traits, bool Is_causal>
void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
    constexpr size_t smem_size = Kernel_traits::kSmemSize;
    // Work-around for gcc 7. It doesn't like nested BOOL_SWITCH.
    // https://github.com/HazyResearch/flash-attention/issues/21
    const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
    dim3 grid(num_m_block, params.b, params.h);
    const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0;
    const bool is_even_K = params.d == Kernel_traits::kHeadDim;
    BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
        BOOL_SWITCH(is_even_K, IsEvenKConst, [&] {
            BOOL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] {
                BOOL_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
                    auto kernel = &flash_fwd_kernel<Kernel_traits, Is_causal, Is_local && !Is_causal, Has_alibi, IsEvenMNConst && IsEvenKConst && !Is_local &&  Kernel_traits::kHeadDim <= 128, IsEvenKConst>;
                    if (smem_size >= 48 * 1024) {
              cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
                    }
                    kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
                });
            });
        });
    });
}

以上接口略过不讨论。讨论适应更复杂情况的 run_flash_splitkv_fwd 接口：

template<typename Kernel_traits>
void run_flash_splitkv_fwd(Flash_fwd_params &params, cudaStream_t stream) {
    static_assert(!Kernel_traits::Is_Q_in_regs, "SplitKV implementation does not support Is_Q_in_regs");
    static_assert(!Kernel_traits::Share_Q_K_smem, "SplitKV implementation does not support Share_Q_K_smem");
    constexpr size_t smem_size = Kernel_traits::kSmemSize;
    const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
    dim3 grid(num_m_block, params.num_splits > 1 ? params.num_splits : params.b, params.num_splits > 1 ? params.b * params.h : params.h);
    const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0;
    const bool is_even_K = params.d == Kernel_traits::kHeadDim;
    BOOL_SWITCH(params.is_causal, Is_causal, [&] {
        BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
            BOOL_SWITCH(is_even_K, IsEvenKConst, [&] {
                BOOL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] {
                    BOOL_SWITCH(params.num_splits > 1, Split, [&] {
                        BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV, [&] {
                            BOOL_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
                                // If Append_KV, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr.
                                // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates.
                                // If Is_local, set Is_causal to false
                                auto kernel = &flash_fwd_splitkv_kernel<Kernel_traits, Is_causal, Is_local && !Is_causal, Has_alibi, IsEvenMNConst && !Append_KV && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128, IsEvenKConst, Split, Append_KV>;
                                if (smem_size >= 48 * 1024) {
                  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
                                }
                                kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
                            });
                        });
                    });
                });
            });
        });
    });
    if (params.num_splits > 1) {
        // We want kBlockM to be as small as possible for more parallelism.
        // With 128 threads we can load 512 elements at a time, so if headdim is divisible by 128, kBlockM = 4.
        // If headdim is divisible by 64, then we set kBlockM = 8, etc.
        constexpr static int kBlockM = Kernel_traits::kHeadDim % 128 == 0 ? 4 : (Kernel_traits::kHeadDim % 64 == 0 ? 8 : 16);
        dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM);
        BOOL_SWITCH(is_even_K, IsEvenKConst, [&] {
            if (params.num_splits <= 2) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 1, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 4) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 2, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 8) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 3, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 16) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 4, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 32) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 5, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 64) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 6, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            } else if (params.num_splits <= 128) {
                flash_fwd_splitkv_combine_kernel<Kernel_traits, kBlockM, 7, IsEvenKConst><<<grid_combine, Kernel_traits::kNThreads, 0, stream>>>(params);
            }
        });
    }
}

首先概括：该接口分为两步，先以 auto kernel = &flash_fwd_splitkv_kernel 和 kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params) 的形式调用处理并保存分块中间结果，之后再调用 flash_fwd_splitkv_combine_kernel 将分块存储的中间结果校正，得到准确的结果。这两个过程对应了《flashatten takeaways EP1 原理》中提到的“onlined stailized softmax” 的分块计算然后逐步合并过程。 Flash_fwd_params &params 由更顶层的接口构建好了之后传入。params中涉及到的以及编译时涉及到的参数在此简单注解：